#' Marginal Effects para felm
#'
#' Funcao para retornar efeitos marginal para modelos felm com interacoes entre tratamento e variavel discreta
#'
#' @param model Um modelo estimado com o pacote lfe
#' @param base_name Um character com o nome do coeficiente do tratamento
#' @param rem_name Um character com o nome da parte dos coeficientes para remover
#' @param ci Intervalo de confianca
#'
#' @importFrom purrr map_dbl
#' @importFrom stats qnorm
#'
#' @export

marg_felm <- function(model, base_name, rem_name, ci = 0.95){
  
  # Extract coefs, cov, and number of coefs
  coefs <- model$coefficients
  ccov <- model$clustervcv
  n <- length(coefs)
  
  # Calculate marginal effects
  marg_coefs <- coefs[1] + coefs[2:n]
  
  # Calculate confidence intervals
  sqr_covs <- purrr::map_dbl(2:n, ~ ccov[.x, .x])
  sqr_main_cov <- 2 * purrr::map_dbl(2:n, ~ ccov[1, .x])
  se <- sqrt(ccov[1, 1] + sqr_covs + sqr_main_cov)
  ci <- stats::qnorm((1 - (1 - ci) / 2))
  ci_up <- marg_coefs + ci * se
  ci_low <- marg_coefs - ci * se
  
  # Inclue baseline coef and cis
  marg_coefs <- c(coefs[1], marg_coefs)
  ci_up <- c(coefs[1] + ci * model$cse[1], ci_up)
  ci_low <- c(coefs[1] - ci * model$cse[1], ci_low)
  
  # Return
  data.frame(vars = c(base_name, gsub(rem_name, "", rownames(coefs)[2:n])),
             coef = marg_coefs, ci_up = ci_up, ci_low = ci_low, stringsAsFactors = FALSE)
}


#' Predict felm
#'
#' Funcao para retornar efeitos preditos de felm
#'
#' @param object Modelo
#' @param newdata Data.frame com dados usados na predicao
#' @param se.fit Retornar intervalos de confianca?
#' @param interval Tipo de intervalo
#' @param level Nivel de confianca
#'
#' @importFrom purrr map_dbl
#' @importFrom stats qnorm
#'
#' @export
predict_felm <- function(object, newdata, se.fit = FALSE,
                         interval = "none",
                         level = 0.95){
  
  tt <- terms(object)
  Terms <- delete.response(tt)
  attr(Terms, "intercept") <- 0
  
  m.mat <- model.matrix(Terms, data = newdata)
  m.coef <- as.numeric(object$coef)
  fit <- as.vector(m.mat[,-2] %*% object$coef)
  fit <- data.frame(fit = fit)
  
  if(se.fit | interval != "none"){
    if(!is.null(object$clustervcv)){
      vcov_mat <- object$clustervcv
    } else if (!is.null(object$robustvcv)) {
      vcov_mat <- object$robustvcv
    } else if (!is.null(object$vcv)){
      vcov_mat <- object$vcv
    } else {
      stop("No vcv attached to felm object.")
    }
    se.fit_mat <- sqrt(diag(m.mat[,-2] %*% vcov_mat %*% t(m.mat[,-2])))
  }
  if(interval == "confidence"){
    t_val <- qt((1 - level) / 2 + level, df = object$df.residual)
    fit$lwr <- fit$fit - t_val * se.fit_mat
    fit$upr <- fit$fit + t_val * se.fit_mat
  } else if (interval == "prediction"){
    stop("interval = \"prediction\" not yet implemented")
  }
  if(se.fit){
    fit$se <- se.fit_mat
    fit <- as.data.frame(cbind(fit, newdata), stringsAsFactors = FALSE)
    return(fit)
  } else {
    return(fit)
  }
}


# Funcao para calcular intervalos preditivos agregados
simula <- function(indicador, pop, fit, se, sims = 1000, unidade = c("bilhoes", "milhoes", "reais")){
  
  unidade <- match.arg(unidade)
  
  if(unidade == "bilhoes"){
    
    mult <- 1000000000
  } else if(unidade == "milhoes"){
    
    mult <- 1000000
  } else {
    
    mult <- 1
  }
  
  distr <- 1:sims %>%
    purrr::map_dbl(~ sum(rnorm(length(indicador), fit, se) * pop, na.rm = T) / mult)
  est <- sum(fit * pop, na.rm = T) / mult
  quants <- quantile(distr, c(0.025, 0.975))
  tibble::tibble(est = est, lwr = quants[1], upr = quants[2])
}


# Funcao para parsear os resultados numa tabela
did_parse <- function(modelo, grupos, media, vars = 1) {
  
  coef <- ifelse(modelo$cpval[vars] < 0.05, 
                 paste0(round(modelo$coefficients[vars], 2),"\\textsuperscript{*}"),
                 as.character(round(modelo$coefficients[vars], 2)))
  data.frame(grupos = grupos,
             coef = coef, 
             se = round(modelo$cse[vars], 2),
             media_controle = round(media, 2),
             N = modelo$N
  )
}